#include "websocket.hpp"

auto http_message::parse_http_message(std::istream& input) -> std::variant<std::string, http_message> {
    static constexpr int eof = std::char_traits<char>::eof();

    http_message message;

    std::string name;
    std::string value;
    name.reserve(32);
    value.reserve(64);

    int ch = input.get();
    while (ch != eof && ch != '\r' && ch != '\n') {
        name.clear();
        value.clear();
        while (ch != eof && ch != ':' && ch != '\n' && name.length() < 256) {
            name += static_cast<char>(ch);
            ch = input.get();
        }
        if (ch == '\n') {
            ch = input.get();
            continue;
        }
        // ignore invalid header lines
        if (ch != ':') return "Field name too long/no colon found";
        if (ch != eof) ch = input.get(); // ':'
        while (ch != eof && ::isspace(ch) && ch != '\r' && ch != '\n') ch = input.get();
        while (ch != eof && ch != '\r' && ch != '\n' && value.length() < 8192) {
            value += static_cast<char>(ch);
            ch = input.get();
        }
        if (ch == '\r') ch = input.get();
        if (ch == '\n') {
            ch = input.get();
        } else if (ch != eof) {
            return "Field value too long/no CRLF found";
        }
        // http header folding
        while (ch == ' ' || ch == '\t') {
            while (ch != eof && ch != '\r' && ch != '\n' && value.length() < 8192) {
                value += static_cast<char>(ch);
                ch = input.get();
            }
            if (ch == '\r') ch = input.get();
            if (ch == '\n') {
                ch = input.get();
            } else if (ch != eof) {
                return "Folded field value too long/no CRLF found";
            }
        }
        size_t v_pos = value.length();
        while (v_pos && ::isspace(value[v_pos - 1])) v_pos--;
        value.resize(v_pos);
        message.key_values[name].push_back(value);
    }

    return message;
}

bool http_message::contains_key_value(const std::string& key, const std::string& value) const {
    auto res = key_values.find(key);
    if (res == key_values.end()) return false;
    auto res2 = std::find(res->second.begin(), res->second.end(), value);
    return res2 != res->second.end();
}

std::string http_message::dump() const {
    std::string result;
    for (const auto& pair: key_values) {
        result += pair.first + ": ";
        result += pair.second[0];
        for (size_t i = 1; i < pair.second.size(); ++i) {
            result += ", " + pair.second[i];
        }
        result += "\r\n";
    }
    return result;
}

auto handshake_request::parse_handshake_request(nonstd::span<char> buf)
    -> std::variant<std::string, handshake_request> {
    static constexpr int eof = std::char_traits<char>::eof();

    handshake_request request;

    std::string method;
    std::string version;
    std::string uri;
    method.reserve(32);
    version.reserve(32);
    uri.reserve(64);

    std::stringstream input(std::string(buf.data(), buf.size()));

    int ch = input.get();
    if (input.bad()) return "Error reading HTTP request header";
    if (ch == eof) return "reached eof";
    while (::isspace(ch)) ch = input.get();
    if (ch == eof) return "No HTTP request header";
    while (!::isspace(ch) && ch != eof && method.length() < 2048) {
        method += static_cast<char>(ch);
        ch = input.get();
    }
    if (!::isspace(ch)) return "HTTP request method invalid or too long";
    while (::isspace(ch)) ch = input.get();
    while (!::isspace(ch) && ch != eof && uri.length() < 2083) {
        uri += static_cast<char>(ch);
        ch = input.get();
    }
    if (!::isspace(ch)) return "HTTP request URI invalid or too long";
    while (::isspace(ch)) ch = input.get();
    while (!::isspace(ch) && ch != eof && version.length() < 2048) {
        version += static_cast<char>(ch);
        ch = input.get();
    }
    if (!::isspace(ch)) return "Invalid HTTP version string";
    while (ch != '\n' && ch != eof) {
        ch = input.get();
    }

    auto _message = http_message::parse_http_message(input);
    if (_message.index() == 0) {
        return std::get<std::string>(_message);
    }

    request.message = std::get<http_message>(_message);

    request.method = method;
    request.version = version;
    request.uri = uri;

    if (request.is_valid_request()) {
        return request;
    } else {
        return "request is not a valid websocket handshake";
    }
}

bool handshake_request::is_valid_request() const {
    if (method != "GET") return false;
    if (!message.contains_key("Host")) return false;
    if (!message.contains_key_value("Connection", "Upgrade")) return false;
    if (!message.contains_key_value("Upgrade", "websocket")) return false;
    if (!message.contains_key("Sec-WebSocket-Key")) return false;
    // if (!contains_key("Sec-WebSocket-Protocol")) return false;
    if (!message.contains_key_value("Sec-WebSocket-Version", "13")) return false;
    // if (!message.contains_key("Origin")) return false;
    return true;
}

auto handshake_response::parse_handshake_response(nonstd::span<char> buf)
    -> std::variant<std::string, handshake_response> {
    static constexpr int eof = std::char_traits<char>::eof();

    handshake_response response;

    std::string version;
    std::string status;
    std::string reason;

    std::stringstream input(std::string(buf.data(), buf.size()));

    int ch = input.get();
    if (input.bad()) return "Error reading HTTP response header";
    if (ch == eof) return "No Message";
    while (::isspace(ch)) ch = input.get();
    if (ch == eof) return "No HTTP response header";
    while (!::isspace(ch) && ch != eof && version.length() < 256) {
        version += static_cast<char>(ch);
        ch = input.get();
    }
    if (!::isspace(ch)) return "Invalid HTTP version string";
    while (::isspace(ch)) ch = input.get();
    while (!::isspace(ch) && ch != eof && status.length() < 256) {
        status += static_cast<char>(ch);
        ch = input.get();
    }
    if (!::isspace(ch)) return "Invalid HTTP status code";
    while (::isspace(ch) && ch != '\r' && ch != '\n' && ch != eof) ch = input.get();
    while (ch != '\r' && ch != '\n' && ch != eof && reason.length() < 2048) {
        reason += static_cast<char>(ch);
        ch = input.get();
    }
    if (!::isspace(ch)) return "HTTP reason string too long";
    if (ch == '\r') ch = input.get();
    if (ch != '\n') return "Unterminated HTTP response line";

    auto _message = http_message::parse_http_message(input);
    if (_message.index() == 0) {
        return std::get<std::string>(_message);
    }

    response.reason = reason;
    response.version = version;
    response.status = status;

    response.message = std::get<http_message>(_message);

    return response;
}

websocket_frame::websocket_frame(websocket_frame&& frame) noexcept {
    FIN = frame.FIN;
    RSV1 = frame.RSV1;
    RSV2 = frame.RSV2;
    RSV3 = frame.RSV3;
    masked = frame.masked;
    opcode = frame.opcode;
    length = frame.length;
    mask_key = frame.mask_key;
    buffer = std::move(frame.buffer);
    frame.FIN = frame.masked = false;
    frame.opcode = CONTINUATION_FRAME;
    frame.length = 0;
    frame.mask_key.fill(0);
}

websocket_frame& websocket_frame::operator=(websocket_frame&& frame) noexcept {
    FIN = frame.FIN;
    RSV1 = frame.RSV1;
    RSV2 = frame.RSV2;
    RSV3 = frame.RSV3;
    masked = frame.masked;
    opcode = frame.opcode;
    length = frame.length;
    mask_key = frame.mask_key;
    buffer = std::move(frame.buffer);
    frame.FIN = frame.masked = false;
    frame.opcode = CONTINUATION_FRAME;
    frame.length = 0;
    frame.mask_key.fill(0);
    return *this;
}

auto websocket_frame::parse_websocket_frame(nonstd::span<char> buf)
    -> std::variant<std::string, size_t, std::pair<websocket_frame, size_t>> {
    const size_t len = buf.size();

    if (len < 2) {
        return 2 - len;
    }

    websocket_frame frame;

    frame.RSV1 = buf[0] & 0b0100'0000;
    frame.RSV2 = buf[0] & 0b0010'0000;
    frame.RSV3 = buf[0] & 0b0001'0000;

    if (frame.RSV1 || frame.RSV2 || frame.RSV3) {
        return "RSV is not all 0 in websocket frame";
    }

    frame.FIN = buf[0] & 0b1000'0000;
    frame.opcode = static_cast<websocket_frame_code>(buf[0] & 0b1111);

    frame.masked = buf[1] & 0b1000'0000;
    uint8_t payload_len = buf[1] & 0b0111'1111;

    size_t offset = 2;
    if (payload_len == 126) {
        if (len < 4) {
            return 4 - len;
        }
        offset += 2;
        frame.length = static_cast<uint8_t>(buf[2]) * 256 + static_cast<uint8_t>(buf[3]);
    } else if (payload_len == 127) {
        if (len < 10) {
            return 10 - len;
        }
        uint64_t length = 0;
        for (int i = 0; i < 8; i++) {
            length = static_cast<uint8_t>(buf[2 + i]) + length * 256ull;
        }
        offset += 8;
        frame.length = length;
    } else {
        frame.length = payload_len;
    }

    if (frame.masked) {
        if (len < offset + 4) {
            return offset + 4 - len;
        }
        std::copy(buf.begin() + offset, buf.begin() + offset + 4, frame.mask_key.begin());
        offset += 4;
    }

    const auto frame_size = static_cast<size_t>(offset + frame.length);

    if (frame_size > len) {
        return frame_size - len;
    }

    frame.buffer = std::vector(buf.begin() + offset, buf.begin() + frame_size);

    if (frame.masked) {
        for (uint64_t i = 0; i < frame.length; i++) {
            frame.buffer[i] ^= frame.mask_key[i % 4];
        }
    }
    return std::make_pair(std::move(frame), frame_size);
}

auto websocket_frame::parse_multiple_websocket_frame(nonstd::span<char> buf)
    -> std::variant<std::string, std::tuple<std::vector<websocket_frame>, size_t, size_t>> {
    std::vector<websocket_frame> frames;

    const auto len = buf.size();

    auto result = websocket_frame::parse_websocket_frame(buf);
    if (result.index() == 0) { // exception
        return std::get<std::string>(result);
    }
    if (result.index() == 1) { // not enough data
        return std::make_tuple(frames, 0, std::get<size_t>(result));
    }
    size_t offset = 0;
    // continue parsing
    while (true) {
        auto [frame, frame_total_size] = std::move(std::get<2>(result));

        if (frame_total_size + offset > len) {
            break;
        }
        offset += static_cast<size_t>(frame_total_size);

        frames.emplace_back(std::move(frame));

        result = websocket_frame::parse_websocket_frame({buf.begin() + offset, len - offset});
        if (result.index() == 0) {
            return std::get<std::string>(result);
        }
        if (result.index() == 1) {
            return std::make_tuple(std::move(frames), offset, std::get<size_t>(result));
        }
    }
    return std::make_tuple(std::move(frames), offset, 0);
}

std::vector<char> websocket_frame::dump() const {
    auto buf = dump_header();
    buf.insert(buf.end(), buffer.begin(), buffer.end());
    return buf;
}

std::vector<char> websocket_frame::dump_header() const {
    std::vector<char> buf(2);
    // write FIN, RSV1, RSV2, RSV3, opcode
    buf[0] |= FIN * 0b1000'0000;
    buf[0] |= RSV1 * 0b0100'0000;
    buf[0] |= RSV2 * 0b0010'0000;
    buf[0] |= RSV3 * 0b0001'0000;
    buf[0] |= opcode;

    // write MASK, Payload len
    buf[1] |= masked * 0b1000'0000;
    if (length > 65535) {
        buf[1] |= 0b0111'1111;
        for (int i = 0; i < 8; i++) {
            buf.push_back(0);
        }
        auto len = length;
        for (int i = 7; i >= 0; i--) {
            buf[2 + i] = len % 256;
            len /= 256;
        }
    } else if (length > 125) {
        buf[1] |= 0b0111'1110;
        for (int i = 0; i < 2; i++) {
            buf.push_back(0);
        }
        auto len = length;
        for (int i = 1; i >= 0; i--) {
            buf[2 + i] = len % 256;
            len /= 256;
        }
    } else {
        buf[1] |= static_cast<char>(length);
    }

    if (masked) {
        for (int i = 0; i < 4; i++) {
            buf.push_back(mask_key[i]);
        }
    }
    return buf;
}

bool websocket_receive_buffer::receive_data(std::vector<char> buffer) {
    std::scoped_lock<std::mutex> lock(_m);

    _buffer.insert(_buffer.end(), buffer.begin(), buffer.end());
    if (buffer.size() < _bytes_required) {
        _bytes_required -= buffer.size();
        return true;
    }

    auto var_result = websocket_frame::parse_multiple_websocket_frame(_buffer);
    if (var_result.index() == 0) {
        lerror("parse websocket frame failed with: " << std::get<std::string>(var_result));
        return false;
    }
    auto& [frames, parsed_size, bytes_required] = std::get<1>(var_result);

    _bytes_required = bytes_required;

    if (parsed_size == 0) {
        return true;
    }

    _buffer.erase(_buffer.begin(), _buffer.begin() + parsed_size);

    // websocket frame should be masked if server
    if (_is_server && std::any_of(frames.begin(), frames.end(), [](const auto& frame) { return !frame.masked; })) {
        lerror("websocket frame sent to server unmasked.");
        return false;
    }

    // filter control frames
    for (auto& frame: frames) {
        if (frame.is_control_frame() && !_receive_hook(std::move(frame))) { // control frame
            return false;
        }
    }

    frames.erase(
        std::remove_if(frames.begin(), frames.end(),
            [](const websocket_frame& frame) { return frame.length == 0 && frame.opcode == CONTINUATION_FRAME; }),
        frames.end());

    // packet still not finish
    if (std::none_of(frames.begin(), frames.end(), [](const auto& frame) { return frame.FIN; })) {
        for (auto& frame: frames) {
            _unhandled_frames.push_back(std::move(frame));
        }
        return true;
    }

    websocket_frame result_frame;
    auto frame_iter = frames.begin();

    // there are remaining data frames
    if (!_unhandled_frames.empty()) {
        result_frame = std::move(_unhandled_frames.front());
        auto total_size = result_frame.buffer.size();
        for (auto iter = _unhandled_frames.begin() + 1; iter != _unhandled_frames.end(); ++iter) {
            total_size += iter->buffer.size();
        }
        result_frame.buffer.reserve(total_size);
        for (auto iter = _unhandled_frames.begin() + 1; iter != _unhandled_frames.end(); ++iter) {
            result_frame.buffer.insert(result_frame.buffer.end(), iter->buffer.begin(), iter->buffer.end());
        }
        _unhandled_frames.clear();
    }

    while (frame_iter != frames.end()) {
        auto fin = std::find_if(frame_iter, frames.end(), [](const auto& frame) { return frame.FIN; });
        if (fin == frames.end()) {
            // no more
            for (auto iter = frame_iter; iter != frames.end(); ++iter) {
                _unhandled_frames.push_back(std::move(*iter));
            }
            break;
        }
        // the first frame should not be continuation frame
        if (result_frame.opcode == CONTINUATION_FRAME) {
            assert(frame_iter->opcode != CONTINUATION_FRAME, "invalid continuatino frame.");
            result_frame = *frame_iter;
            ++frame_iter;
        }
        auto total_size = result_frame.buffer.size();
        for (auto iter = frame_iter; iter != fin + 1; ++iter) {
            total_size += iter->buffer.size();
        }
        result_frame.buffer.reserve(total_size);
        for (auto iter = frame_iter; iter != fin + 1; ++iter) {
            result_frame.buffer.insert(result_frame.buffer.end(), iter->buffer.begin(), iter->buffer.end());
        }
        result_frame.length = result_frame.buffer.size();
        result_frame.FIN = true;
        if (!_receive_hook(std::move(result_frame))) {
            return false;
        }
        result_frame = websocket_frame();
        frame_iter = fin + 1;
    }
    return true;
}
